﻿using System.Collections.Generic;
using UnityEngine;
using UnityEditor;

namespace CWM.Skinn
{
    public static partial class SkinnContextCommands
    {
        [CtxMenuItem("Renderer/Export/Experimental/Atlas/512", "Combine Submeshes and Render.", ContextFlag.None, ContextRequirements.SubMeshes, 0)]
        public static void RendererCombineAtlasExport512() { RendererCombineAtlas(512); }

        [CtxMenuItem("Renderer/Export/Experimental/Atlas/1024", "Combine Submeshes and Render.", ContextFlag.None, ContextRequirements.SubMeshes, 0)]
        public static void RendererCombineAtlasExport1024() { RendererCombineAtlas(1024); }

        [CtxMenuItem("Renderer/Export/Experimental/Atlas/2048", "Combine Submeshes and Render.", ContextFlag.None, ContextRequirements.SubMeshes, 0)]
        public static void RendererCombineAtlasExport2048() { RendererCombineAtlas(2048); }

        [CtxMenuItem("Renderer/Export/Experimental/Atlas/4096", "Combine Submeshes and Render.", ContextFlag.None, ContextRequirements.SubMeshes, 0)]
        public static void RendererCombineAtlasExport4096() { RendererCombineAtlas(4096); }

        [CtxMenuItem("Renderer/Export/Experimental/Atlas HardSurface/2048", "Combine Submeshes and Render.", ContextFlag.None, ContextRequirements.SubMeshes, 0)]
        public static void RendererCombineAtlasExportHardSurface2048() { RendererCombineAtlas(2048, MeshSurfaceType.HardSuface); }

        public static void RendererCombineAtlas(int size, MeshSurfaceType meshSurfaceType = MeshSurfaceType.Orangic)
        {
            switch (meshSurfaceType)
            {
                case MeshSurfaceType.Orangic: LMAtlasUtility.DefualtOrangicSettings();break;
                case MeshSurfaceType.HardSuface: LMAtlasUtility.HardSurfceSettings(); break;
                case MeshSurfaceType.Mixed: LMAtlasUtility.MixedSettings(); break;
                default: break;
            }

            foreach (var item in SkinnContext.GetContextItems())
            {
                var context = item.GetComponent<Renderer>();
                if (SkinnEx.IsNullOrNotInAScene(context)) continue;
                if (SkinnEx.IsNullOrNotVailid(context)) continue;
                if(!SkinnContext.IsVertexCountOkay(context)) continue;
                LMAtlasUtility.ExportAtlas(context, size);
            }
            AssetDatabase.Refresh(ImportAssetOptions.ForceUpdate);

            SkinnContext.Finished();
        }

        public static class LMAtlasUtility
        {
            public static int UVHardAngle = 180;//0-180 default 90
            public static int UVPackMargin = 2; //1-64 default
            public static int UVAngleDistortion = 2; //1-100 default 8
            public static int UVAreaDistortion = 2; //1-100 default 15

            public static void DefualtOrangicSettings()
            {
                UVHardAngle = 180; UVPackMargin = 2; UVAngleDistortion = 2; UVAreaDistortion = 2;
            }

            public static void MixedSettings()
            {
                UVHardAngle = 90; UVPackMargin = 4; UVAngleDistortion = 8; UVAreaDistortion = 15;
            }

            public static void HardSurfceSettings()
            {
                UVHardAngle = 60; UVPackMargin = 4; UVAngleDistortion = 8; UVAreaDistortion = 15;
            }

            public static void ExportAtlas(Renderer context, int atlasSize)
            {
                if (SkinnEx.IsNullOrNotVailid(context)) return;
                var assetPath = EditorUtility.SaveFolderPanel("Export Atlas FBX", "Assets/", "");
                if (string.IsNullOrEmpty(assetPath)) return;
                if (assetPath.StartsWith(Application.dataPath)) assetPath = "Assets" + assetPath.Substring(Application.dataPath.Length);
                var contextName = context.name + "_Combined";
                var rootFolder = assetPath + "/" + contextName;
                if (!AssetDatabase.IsValidFolder(rootFolder)) AssetDatabase.CreateFolder(assetPath, contextName);
                var imageFolder = rootFolder + "/" + "Materials";
                if (!AssetDatabase.IsValidFolder(imageFolder)) AssetDatabase.CreateFolder(rootFolder, "Materials");

                var contextMesh = context.GetSharedMesh();
                var tempMesh = contextMesh.Clone() as Mesh;
                var tempUVA = tempMesh.uv2;
                var tempUVB = tempMesh.uv3;
                if (!SkinnEx.IsNullOrEmpty(tempUVA)) tempMesh.uv3 = tempUVA;
                if (!SkinnEx.IsNullOrEmpty(tempUVB)) tempMesh.uv4 = tempUVB;
                tempMesh.uv2 = null;
                context.SetSharedMesh(tempMesh);

                var tempPath = string.Format("Assets/TempFBX_{0}.fbx", context.name);
                string fbx;

                if (context as SkinnedMeshRenderer) fbx = (context as SkinnedMeshRenderer).ToACSIIFBX(true);
                else fbx = (context as MeshRenderer).ToACSIIFBX();

                if (string.IsNullOrEmpty(fbx)) return;
                if (System.IO.File.Exists(tempPath)) AssetDatabase.DeleteAsset(tempPath);
                System.IO.File.WriteAllText(tempPath, fbx);

                SkinnEx.Release(tempMesh); tempUVA = null; tempUVB = null;
                context.SetSharedMesh(contextMesh);

                AssetDatabase.Refresh();
                var importer = ModelImporter.GetAtPath(tempPath) as ModelImporter;

                importer.animationType = ModelImporterAnimationType.Generic;
                importer.importNormals = ModelImporterNormals.Import;
                importer.normalCalculationMode = ModelImporterNormalCalculationMode.Unweighted;
                importer.normalSmoothingAngle = 60f;
                importer.importTangents = ModelImporterTangents.None;
                importer.importBlendShapes = true;
                importer.meshCompression = ModelImporterMeshCompression.Off;
                importer.weldVertices = false;
                importer.materialName = ModelImporterMaterialName.BasedOnMaterialName;

                importer.generateSecondaryUV = true;
                importer.secondaryUVHardAngle = UVHardAngle;
                importer.secondaryUVPackMargin = UVPackMargin;
                importer.secondaryUVAngleDistortion = UVAngleDistortion;
                importer.secondaryUVAreaDistortion = UVAreaDistortion;

                AssetDatabase.WriteImportSettingsIfDirty(tempPath);
                AssetDatabase.ImportAsset(tempPath, ImportAssetOptions.ForceUncompressedImport);

                var export = AssetDatabase.LoadAssetAtPath(tempPath, typeof(GameObject)) as GameObject;
                if (!export) { Debug.LogError("Failed To Import"); return; }
                var tempRenderer = GameObject.Instantiate(export).GetComponentInChildren<Renderer>();
                tempRenderer.name = context.name;
                tempMesh = tempRenderer.GetSharedMesh().Clone() as Mesh;

                var sharedMaterials = context.GetSharedMaterialsFromSubMeshes();
                tempRenderer.sharedMaterials = sharedMaterials;
                var textures = RenderAllTextures(tempRenderer, imageFolder + "/", atlasSize);
                tempMesh.subMeshCount = 1;
                tempMesh.SetTriangles(tempRenderer.GetSharedMesh().triangles, 0);

                var maxMaterialIndex = -1;
                var maxPropertyCount = -1;
                for (int i = 0; i < sharedMaterials.Length; i++)
                {
                    var propertyCount = 0;
                    if (!sharedMaterials[i]) continue;
                    foreach (var item in textures) if (sharedMaterials[i].HasProperty(item.Key)) propertyCount++;
                    if (propertyCount > maxPropertyCount) maxMaterialIndex = i;
                }

                if (maxMaterialIndex > -1)
                {
                    AssetDatabase.Refresh();

                    var combinedMat = new Material(sharedMaterials[maxMaterialIndex].shader) { name = contextName };
                    foreach (var item in textures)
                    {
                        var texture = (Texture)AssetDatabase.LoadAssetAtPath(item.Value, typeof(Texture));
                        if (!texture) continue;
                        if (combinedMat.GetColorProperty() == item.Key)
                        {
                            if (!textures.ContainsKey("_MainTex") && combinedMat.HasProperty("_MainTex")) combinedMat.SetTexture("_MainTex", texture);
                        }
                        else if (combinedMat.HasProperty(item.Key)) combinedMat.SetTexture(item.Key, texture);

                        var textureImporter = TextureImporter.GetAtPath(item.Value) as TextureImporter;
                        if (!textureImporter) continue;

                        if (item.Key.Contains("Bump") || item.Key.Contains("Normal"))
                            textureImporter.textureType = TextureImporterType.NormalMap;
                        textureImporter.alphaIsTransparency = true;

                        textureImporter.mipmapEnabled = false;

                        AssetDatabase.WriteImportSettingsIfDirty(item.Value);
                        AssetDatabase.ImportAsset(item.Value, ImportAssetOptions.ForceUncompressedImport);
                    }
                    tempRenderer.sharedMaterials = new Material[] { combinedMat };
                    AssetDatabase.CreateAsset(combinedMat, imageFolder + "/" + combinedMat.name + ".mat");
                }
                else tempRenderer.sharedMaterials = new Material[1];

                tempUVA = tempMesh.uv2;
                tempUVB = tempMesh.uv;
                tempMesh.uv = tempUVA;
                tempMesh.uv2 = tempUVB;

                tempRenderer.SetSharedMesh(tempMesh);

                if (tempRenderer as SkinnedMeshRenderer) fbx = (tempRenderer as SkinnedMeshRenderer).ToACSIIFBX(true);
                else fbx = (tempRenderer as MeshRenderer).ToACSIIFBX();

                var fbxPath = rootFolder + "/" + contextName + ".fbx";
                if (System.IO.File.Exists(fbxPath)) AssetDatabase.DeleteAsset(fbxPath);
                System.IO.File.WriteAllText(fbxPath, fbx);

                AssetDatabase.Refresh();
                importer = ModelImporter.GetAtPath(fbxPath) as ModelImporter;

                if (context as SkinnedMeshRenderer) importer.animationType = ModelImporterAnimationType.Generic;
                else importer.animationType = ModelImporterAnimationType.None;

                importer.importNormals = ModelImporterNormals.Import;
                importer.normalCalculationMode = ModelImporterNormalCalculationMode.Unweighted;
                importer.normalSmoothingAngle = 60f;
                importer.importTangents = ModelImporterTangents.CalculateLegacyWithSplitTangents;
                importer.meshCompression = ModelImporterMeshCompression.Off;
                importer.weldVertices = false;
                importer.materialName = ModelImporterMaterialName.BasedOnMaterialName;
                importer.materialLocation = ModelImporterMaterialLocation.InPrefab;
                importer.materialSearch = ModelImporterMaterialSearch.Local;
                importer.generateSecondaryUV = false;
                importer.secondaryUVHardAngle = UVHardAngle;
                importer.secondaryUVPackMargin = UVPackMargin;
                importer.secondaryUVAngleDistortion = UVAngleDistortion;
                importer.secondaryUVAreaDistortion = UVAreaDistortion;
                importer.SearchAndRemapMaterials(ModelImporterMaterialName.BasedOnMaterialName, ModelImporterMaterialSearch.Local);

                AssetDatabase.Refresh();

                var finalExport = AssetDatabase.LoadAssetAtPath(fbxPath, typeof(GameObject)) as GameObject;
                //PrefabUtility.InstantiatePrefab(finalExport);
                Selection.objects = new Object[] { finalExport };

                SkinnEx.Release(tempMesh);
                SkinnEx.Release(tempRenderer.transform.root.gameObject);
                AssetDatabase.DeleteAsset(tempPath);
            }

            public static Dictionary<string, string> RenderAllTextures(Renderer context, string path, int atlasSize)
            {
                const bool padTextures = true;
                var textures = new Dictionary<string, string>();

                if (!SkinnEx.HasMinimumRequirements(context)) return textures;
                var shader = SkinnInternalAsset.Asset.lightMapAtlas;
                var renderMaterial = new Material(shader) { name = "uv", hideFlags = HideFlags.HideAndDontSave };
                var mesh = context.GetSharedMesh();
                var materials = context.GetSharedMaterialsFromSubMeshes();
                var matDictionary = new Dictionary<string, List<int>>();
                for (int i = 0; i < materials.Length; i++)
                {
                    var material = materials[i]; if (!material) continue;
                    string[] textureIDs = material.GetAllPropertiesWithTextures();
                    for (int ii = 0; ii < textureIDs.Length; ii++)
                    {
                        var property = textureIDs[ii];
                        if (matDictionary.ContainsKey(property))
                        {
                            List<int> mats = new List<int>();
                            matDictionary.TryGetValue(property, out mats);
                            matDictionary.Remove(property);
                            mats.Add(i);
                            matDictionary.Add(property, mats);
                        }
                        else matDictionary.Add(property, new List<int>(1) { i });
                    }

                    if (textureIDs.Length > 0) continue;
                    var colorProp = material.GetColorProperty();
                    if (string.IsNullOrEmpty(colorProp)) continue;

                    if (matDictionary.ContainsKey(colorProp))
                    {
                        List<int> mats = new List<int>();
                        matDictionary.TryGetValue(colorProp, out mats);
                        matDictionary.Remove(colorProp);
                        mats.Add(i);
                        matDictionary.Add(colorProp, mats);
                    }
                    else matDictionary.Add(colorProp, new List<int>(1) { i });
                }

                foreach (KeyValuePair<string, List<int>> mat in matDictionary)
                {

                    var materialIndex = mat.Value;
                    string assetPath = path + context.name + "_" + mat.Key + "_Atlas.png";
                    float isNormalMap = mat.Key.Contains("Bump") || mat.Key.Contains("Normal") ? 1 : 0;
                    renderMaterial.SetFloat("_ConvertBumpMap", isNormalMap);

                    var atlasRTexture = new RenderTexture(atlasSize, atlasSize, 0, RenderTextureFormat.ARGB32, RenderTextureReadWrite.sRGB) { };

                    var tempRTexture = new RenderTexture(atlasSize, atlasSize, 0) { enableRandomWrite = true };
                    tempRTexture.Create();

                    var active = RenderTexture.active;
                    RenderTexture.active = atlasRTexture;
                    {
                        for (int i = 0; i < materialIndex.Count; i++)
                        {
                            var material = materials[materialIndex[i]];

                            if (material.GetColorProperty() == mat.Key)
                            {
                                renderMaterial.SetFloat("_UseColor", 1);

                                Color propColor = material.GetColor(mat.Key);
                                propColor.a = Mathf.Clamp(propColor.a, 1e-8f, 1);
                                renderMaterial.SetColor("_Color", propColor);
                                renderMaterial.SetTexture("_MainTex", null);

                            }
                            else
                            {
                                var offest = material.GetTextureOffset(mat.Key);
                                var scale = material.GetTextureScale(mat.Key);

                                renderMaterial.SetTextureOffset("_MainTex", offest);
                                renderMaterial.SetTextureOffset("_MainTex", scale);

                                if (mat.Key == "_MainTex" && material.HasProperty("_Color"))
                                {
                                    renderMaterial.SetFloat("_UseColor", 1);
                                    renderMaterial.SetColor("_Color", material.GetColor("_Color"));
                                    renderMaterial.SetTexture("_MainTex", material.GetTexture(mat.Key));
                                }
                                else if (mat.Key == "_EmissionMap" && material.HasProperty("_EmissionColor"))
                                {
                                    renderMaterial.SetFloat("_UseColor", 1);
                                    var emissionColor = material.GetColor("_EmissionColor");
                                    emissionColor.a = 1;
                                    renderMaterial.SetColor("_Color", emissionColor);
                                    renderMaterial.SetTexture("_MainTex", material.GetTexture(mat.Key));
                                }
                                else
                                {
                                    renderMaterial.SetFloat("_UseColor", 0);
                                    renderMaterial.SetColor("_Color", Color.white);
                                    renderMaterial.SetTexture("_MainTex", material.GetTexture(mat.Key));
                                }
                                //TODO// Maps; Spec, Metal, Bump, Detail, Opacity.
                            }

                            renderMaterial.SetPass(0);

                            Graphics.DrawMeshNow(mesh, Matrix4x4.identity, materialIndex[i]);
                            Graphics.DrawMeshNow(mesh, Matrix4x4.identity, materialIndex[i]);

                        }

                        if (padTextures) for (int i = 0; i < 3; i++) PadTexture(atlasRTexture, tempRTexture, 1.001f + (i * 0.0001f));

                        var format = TextureFormat.ARGB32;
                        var newTexture = new Texture2D(atlasSize, atlasSize, format, false, true);
                        newTexture.ReadPixels(new Rect(0, 0, atlasSize, atlasSize), 0, 0);
                        newTexture.Apply(false, false);
                        var bytes = newTexture.EncodeToPNG();
                        if (System.IO.File.Exists(assetPath)) AssetDatabase.DeleteAsset(assetPath);
                        System.IO.File.WriteAllBytes(assetPath, bytes);

                        SkinnEx.Release(newTexture);
                    }

                    RenderTexture.active = active;
                    SkinnEx.Release(atlasRTexture);
                    SkinnEx.Release(tempRTexture);

                    textures.Add(mat.Key, assetPath);
                }

                SkinnEx.Release(renderMaterial);

                if (matDictionary.Count < 1) SkinnContext.LogWarning("No Textures Found", true);

                DefualtOrangicSettings();
                return textures;
            }

            private static void PadTexture(RenderTexture source, RenderTexture temp, float size = 1.1f)
            {
                ComputeShader shader = SkinnInternalAsset.Asset.texturePadding;
                var kernal = shader.FindKernel("AlphaPadding");
                shader.SetTexture(kernal, "Source", source);
                shader.SetTexture(kernal, "Destination", temp);
                shader.SetFloat("Size", size);
                shader.Dispatch(kernal, (temp.width + 7) / 8, (temp.height + 7) / 8, 1);
                Graphics.Blit(temp, source);
            }

            private static void PadTextureEdgeBlending(RenderTexture source, RenderTexture temp, float size = 1.1f)
            {
                ComputeShader shader = SkinnInternalAsset.Asset.texturePadding;
                var kernal = shader.FindKernel("EdgeBlending");
                shader.SetTexture(kernal, "Source", source);
                shader.SetTexture(kernal, "Destination", temp);
                shader.SetFloat("Size", size);
                shader.Dispatch(kernal, (temp.width + 7) / 8, (temp.height + 7) / 8, 1);
                Graphics.Blit(temp, source);
            }
        }
    }
}
